from ollama import chat
from ollama import ChatResponse
from ollama import Client

class BaseModelHandler:
    """模型处理器基类"""
    def __init__(self, window):
        self.window = window
    
    def handle(self, query, model_config):
        """处理请求的抽象方法"""
        raise NotImplementedError

class BlPtHandler(BaseModelHandler):
    """BlPt模型处理器"""
    def handle(self, query, model_config):
        # GPT-4特定的处理逻辑
        self.window.evaluate_js(f"updateMsg('BlPtHandler')")
        return f"BlPtHandler处理结果: {query}"

class DiffOllamaDkr1Handler(BaseModelHandler):
    """DiffOllamaDkr1模型处理器"""
    def handle(self, query, model_config):
        # Claude特定的处理逻辑
        self.window.evaluate_js(f"updateMsg('DiffOllamaDkr1Handler')")
        return f"DiffOllamaDkr1Handler处理结果: {query}"

class OllamaDkr1Handler(BaseModelHandler):
    """OllamaDkr1模型处理器"""
    def handle(self, query, model_config):
        print(model_config)
        model = model_config["params"]['model']
        client = Client(
            host = model_config['endpoint'],
            headers={'x-some-header': 'some-value'}
        )
        stream = client.chat(
            model = model,
            #'deepseek-r1:7b',
            messages = [{'role': 'user', 'content': query}],
            stream = True,
        )
        for chunk in stream:
            msg = chunk['message']['content']
            # print(msg, end='', flush=True)
            if msg:
                self.window.evaluate_js(f"updateMsg('{msg}')")
           
        return f"OllamaDkr1Handler处理结果: {query}"

import os
from cozepy import COZE_CN_BASE_URL
from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageContentType, ChatEventType  # noqa

coze_api_token = 'pat_RzzyCSqEWNH6cdflYtJ2AKxCLzVXiEPw6pq3yH2AhyUdWTTH7XrPih7BO0T0y1V5'
# coze_api_base = COZE_CN_BASE_URL
# coze = Coze(auth=TokenAuth(token=coze_api_token), base_url=coze_api_base)
# bot_id = '7489081421715243023'
# user_id = '389926'

class KouZiHanler(BaseModelHandler):
    def send_command (self, command):
        result = self.window.evaluate_js(f"executeTerminalCommandForAi(\"{command}\")")
        return f"命令执行结果: {result}"


    """KouZiHanler模型处理器"""
    def handle(self, query, model_config):
        # print(model_config)
        bot_id=model_config["params"]['bot_id']
        user_id=model_config["params"]['user_id']
        apikey = model_config["apiKey"]
        coze = Coze(auth=TokenAuth(token=apikey), base_url=COZE_CN_BASE_URL)
        print("COZE_CN_BASE_URL:"+COZE_CN_BASE_URL)
        # 收集所有消息内容
        full_response = ""
        for event in coze.chat.stream(
            bot_id=bot_id,
            user_id=user_id,
            additional_messages=[
                Message.build_user_question_text(query)
            ],
        ):
            if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                print(event.message.content, end="", flush=True)
                msg = event.message.content
                full_response += msg
                if msg:
                    self.window.evaluate_js(f"updateMsg('{msg}')")
            if event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                print("token usage:", event.chat.usage.token_count)

        # 尝试解析JSON格式的响应
        try:
            import json
            import re
            from operator import methodcaller
            
            # 提取JSON部分
            json_match = re.search(r'```json\n(.*?)\n```', full_response, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
                tool_data = json.loads(json_str)
                
                if 'tool_calls' in tool_data:
                    for tool_call in tool_data['tool_calls']:
                        if 'tool_call' in tool_call:
                            tName = tool_call['tool_call']['tool_name']
                            tArgs = tool_call['tool_call']['tool_args']
                            
                            print(f"开始调用工具: {tName} 参数: {tArgs}")
                            caller = methodcaller(tName, **tArgs)
                            result = caller(self)  # 假设工具方法是当前类的成员
                            print(f"工具调用结果: {result}")
                            
                            # 将结果返回给前端
                            self.window.evaluate_js(f"updateMsg('\\n工具调用结果: {result}')")
                            return f"工具调用完成: {tName}"
        except Exception as e:
            print(f"解析或执行工具调用失败: {e}")
        
        return full_response






    